{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Copyright 2021-2022 @ Shenzhen Bay Laboratory & Peking University & Huawei Technologies Co., Ltd\n",
    "\n",
    "This code is a part of Cybertron package.\n",
    "\n",
    "The Cybertron is open-source software based on the AI-framework:\n",
    "MindSpore (https://www.mindspore.cn/)\n",
    "\n",
    "Licensed under the Apache License, Version 2.0 (the \"License\");\n",
    "you may not use this file except in compliance with the License.\n",
    "\n",
    "You may obtain a copy of the License at http://www.apache.org/licenses/LICENSE-2.0\n",
    "\n",
    "Unless required by applicable law or agreed to in writing, software\n",
    "distributed under the License is distributed on an \"AS IS\" BASIS,\n",
    "WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
    "\n",
    "See the License for the specific language governing permissions and\n",
    "limitations under the License.\n",
    "\n",
    "Cybertron tutorial 05: Multi-task with multiple readouts (example 1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "[WARNING] ME(7681:140517018826560,MainProcess):2022-08-10-17:16:35.608.277 [mindspore/run_check/_check_version.py:137] Can not found cuda libs, please confirm that the correct cuda version has been installed, you can refer to the installation guidelines: https://www.mindspore.cn/install\n"
     ]
    }
   ],
   "source": [
    "import sys\n",
    "import time\n",
    "import numpy as np\n",
    "import mindspore as ms\n",
    "from mindspore import nn\n",
    "from mindspore import Tensor\n",
    "from mindspore import context\n",
    "from mindspore import dataset as ds\n",
    "from mindspore.train import Model\n",
    "from mindspore.train.callback import ModelCheckpoint, CheckpointConfig\n",
    "\n",
    "from cybertron import Cybertron\n",
    "from cybertron import MolCT\n",
    "from cybertron import AtomwiseReadout\n",
    "from cybertron.train import MAE, MLoss\n",
    "from cybertron.train import WithLabelLossCell, WithLabelEvalCell\n",
    "from cybertron.train import TrainMonitor\n",
    "from cybertron.train import TransformerLR\n",
    "\n",
    "context.set_context(mode=context.GRAPH_MODE, device_target=\"GPU\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "data_name = sys.path[0] + '/dataset_qm9_normed_'\n",
    "train_file = data_name + 'trainset_1024.npz'\n",
    "valid_file = data_name + 'validset_128.npz'\n",
    "\n",
    "train_data = np.load(train_file)\n",
    "valid_data = np.load(valid_file)\n",
    "\n",
    "idx = [7, 8, 9, 10]  # U0,U,G,H\n",
    "\n",
    "num_atom = int(train_data['num_atoms'])\n",
    "scale = Tensor(train_data['scale'][idx], ms.float32)\n",
    "shift = Tensor(train_data['shift'][idx], ms.float32)\n",
    "ref = Tensor(train_data['type_ref'][:, idx], ms.float32)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "mod = MolCT(\n",
    "    cutoff=1,\n",
    "    n_interaction=3,\n",
    "    dim_feature=128,\n",
    "    n_heads=8,\n",
    "    activation='swish',\n",
    "    max_cycles=1,\n",
    "    length_unit='nm',\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "readout0 = AtomwiseReadout(mod, dim_output=1)\n",
    "readout1 = AtomwiseReadout(mod, dim_output=1)\n",
    "readout2 = AtomwiseReadout(mod, dim_output=1)\n",
    "readout3 = AtomwiseReadout(mod, dim_output=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "net = Cybertron(mod, readout=[\n",
    "    readout0, readout1, readout2, readout3], num_atoms=num_atom, length_unit='nm')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Cybertron<\n",
       "  (model): MolCT<\n",
       "    (activation): Swish<\n",
       "      (sigmoid): Sigmoid<>\n",
       "      >\n",
       "    (atom_embedding): Embedding<vocab_size=64, embedding_size=128, use_one_hot=True, embedding_table=Parameter (name=model.atom_embedding.embedding_table, shape=(64, 128), dtype=Float32, requires_grad=True), dtype=Float32, padding_idx=None>\n",
       "    (cutoff_fn): SmoothCutoff<>\n",
       "    (rbf): LogGaussianBasis<>\n",
       "    (dis_filter): ResFilter<\n",
       "      (linear): Dense<input_channels=64, output_channels=128, has_bias=True>\n",
       "      (residual): Residual<\n",
       "        (nonlinear): MLP<\n",
       "          (mlp): SequentialCell<\n",
       "            (0): Dense<\n",
       "              input_channels=128, output_channels=128, has_bias=True, activation=Swish<>\n",
       "              (activation): Swish<\n",
       "                (sigmoid): Sigmoid<>\n",
       "                >\n",
       "              >\n",
       "            (1): Dense<input_channels=128, output_channels=128, has_bias=True>\n",
       "            >\n",
       "          >\n",
       "        >\n",
       "      >\n",
       "    (interactions): CellList<\n",
       "      (0): NeuralInteractionUnit<\n",
       "        (activation): Swish<\n",
       "          (sigmoid): Sigmoid<>\n",
       "          >\n",
       "        (dis_filter): ResFilter<\n",
       "          (linear): Dense<input_channels=64, output_channels=128, has_bias=True>\n",
       "          (residual): Residual<\n",
       "            (nonlinear): MLP<\n",
       "              (mlp): SequentialCell<\n",
       "                (0): Dense<\n",
       "                  input_channels=128, output_channels=128, has_bias=True, activation=Swish<>\n",
       "                  (activation): Swish<\n",
       "                    (sigmoid): Sigmoid<>\n",
       "                    >\n",
       "                  >\n",
       "                (1): Dense<input_channels=128, output_channels=128, has_bias=True>\n",
       "                >\n",
       "              >\n",
       "            >\n",
       "          >\n",
       "        (positional_embedding): PositionalEmbedding<\n",
       "          (norm): LayerNorm<normalized_shape=(128,), begin_norm_axis=-1, begin_params_axis=-1, gammaParameter (name=model.interactions.0.positional_embedding.norm.gamma, shape=(128,), dtype=Float32, requires_grad=True), beta=Parameter (name=model.interactions.0.positional_embedding.norm.beta, shape=(128,), dtype=Float32, requires_grad=True)>\n",
       "          (x_norm): LayerNorm<normalized_shape=(128,), begin_norm_axis=-1, begin_params_axis=-1, gammaParameter (name=model.interactions.0.positional_embedding.norm.gamma, shape=(128,), dtype=Float32, requires_grad=True), beta=Parameter (name=model.interactions.0.positional_embedding.norm.beta, shape=(128,), dtype=Float32, requires_grad=True)>\n",
       "          (g_norm): LayerNorm<normalized_shape=(128,), begin_norm_axis=-1, begin_params_axis=-1, gammaParameter (name=model.interactions.0.positional_embedding.norm.gamma, shape=(128,), dtype=Float32, requires_grad=True), beta=Parameter (name=model.interactions.0.positional_embedding.norm.beta, shape=(128,), dtype=Float32, requires_grad=True)>\n",
       "          (x2q): Dense<input_channels=128, output_channels=128>\n",
       "          (x2k): Dense<input_channels=128, output_channels=128>\n",
       "          (x2v): Dense<input_channels=128, output_channels=128>\n",
       "          >\n",
       "        (multi_head_attention): MultiheadAttention<\n",
       "          (output): Dense<input_channels=128, output_channels=128>\n",
       "          (softmax_with_mask): SoftmaxWithMask<>\n",
       "          >\n",
       "        >\n",
       "      (1): NeuralInteractionUnit<\n",
       "        (activation): Swish<\n",
       "          (sigmoid): Sigmoid<>\n",
       "          >\n",
       "        (dis_filter): ResFilter<\n",
       "          (linear): Dense<input_channels=64, output_channels=128, has_bias=True>\n",
       "          (residual): Residual<\n",
       "            (nonlinear): MLP<\n",
       "              (mlp): SequentialCell<\n",
       "                (0): Dense<\n",
       "                  input_channels=128, output_channels=128, has_bias=True, activation=Swish<>\n",
       "                  (activation): Swish<\n",
       "                    (sigmoid): Sigmoid<>\n",
       "                    >\n",
       "                  >\n",
       "                (1): Dense<input_channels=128, output_channels=128, has_bias=True>\n",
       "                >\n",
       "              >\n",
       "            >\n",
       "          >\n",
       "        (positional_embedding): PositionalEmbedding<\n",
       "          (norm): LayerNorm<normalized_shape=(128,), begin_norm_axis=-1, begin_params_axis=-1, gammaParameter (name=model.interactions.1.positional_embedding.norm.gamma, shape=(128,), dtype=Float32, requires_grad=True), beta=Parameter (name=model.interactions.1.positional_embedding.norm.beta, shape=(128,), dtype=Float32, requires_grad=True)>\n",
       "          (x_norm): LayerNorm<normalized_shape=(128,), begin_norm_axis=-1, begin_params_axis=-1, gammaParameter (name=model.interactions.1.positional_embedding.norm.gamma, shape=(128,), dtype=Float32, requires_grad=True), beta=Parameter (name=model.interactions.1.positional_embedding.norm.beta, shape=(128,), dtype=Float32, requires_grad=True)>\n",
       "          (g_norm): LayerNorm<normalized_shape=(128,), begin_norm_axis=-1, begin_params_axis=-1, gammaParameter (name=model.interactions.1.positional_embedding.norm.gamma, shape=(128,), dtype=Float32, requires_grad=True), beta=Parameter (name=model.interactions.1.positional_embedding.norm.beta, shape=(128,), dtype=Float32, requires_grad=True)>\n",
       "          (x2q): Dense<input_channels=128, output_channels=128>\n",
       "          (x2k): Dense<input_channels=128, output_channels=128>\n",
       "          (x2v): Dense<input_channels=128, output_channels=128>\n",
       "          >\n",
       "        (multi_head_attention): MultiheadAttention<\n",
       "          (output): Dense<input_channels=128, output_channels=128>\n",
       "          (softmax_with_mask): SoftmaxWithMask<>\n",
       "          >\n",
       "        >\n",
       "      (2): NeuralInteractionUnit<\n",
       "        (activation): Swish<\n",
       "          (sigmoid): Sigmoid<>\n",
       "          >\n",
       "        (dis_filter): ResFilter<\n",
       "          (linear): Dense<input_channels=64, output_channels=128, has_bias=True>\n",
       "          (residual): Residual<\n",
       "            (nonlinear): MLP<\n",
       "              (mlp): SequentialCell<\n",
       "                (0): Dense<\n",
       "                  input_channels=128, output_channels=128, has_bias=True, activation=Swish<>\n",
       "                  (activation): Swish<\n",
       "                    (sigmoid): Sigmoid<>\n",
       "                    >\n",
       "                  >\n",
       "                (1): Dense<input_channels=128, output_channels=128, has_bias=True>\n",
       "                >\n",
       "              >\n",
       "            >\n",
       "          >\n",
       "        (positional_embedding): PositionalEmbedding<\n",
       "          (norm): LayerNorm<normalized_shape=(128,), begin_norm_axis=-1, begin_params_axis=-1, gammaParameter (name=model.interactions.2.positional_embedding.norm.gamma, shape=(128,), dtype=Float32, requires_grad=True), beta=Parameter (name=model.interactions.2.positional_embedding.norm.beta, shape=(128,), dtype=Float32, requires_grad=True)>\n",
       "          (x_norm): LayerNorm<normalized_shape=(128,), begin_norm_axis=-1, begin_params_axis=-1, gammaParameter (name=model.interactions.2.positional_embedding.norm.gamma, shape=(128,), dtype=Float32, requires_grad=True), beta=Parameter (name=model.interactions.2.positional_embedding.norm.beta, shape=(128,), dtype=Float32, requires_grad=True)>\n",
       "          (g_norm): LayerNorm<normalized_shape=(128,), begin_norm_axis=-1, begin_params_axis=-1, gammaParameter (name=model.interactions.2.positional_embedding.norm.gamma, shape=(128,), dtype=Float32, requires_grad=True), beta=Parameter (name=model.interactions.2.positional_embedding.norm.beta, shape=(128,), dtype=Float32, requires_grad=True)>\n",
       "          (x2q): Dense<input_channels=128, output_channels=128>\n",
       "          (x2k): Dense<input_channels=128, output_channels=128>\n",
       "          (x2v): Dense<input_channels=128, output_channels=128>\n",
       "          >\n",
       "        (multi_head_attention): MultiheadAttention<\n",
       "          (output): Dense<input_channels=128, output_channels=128>\n",
       "          (softmax_with_mask): SoftmaxWithMask<>\n",
       "          >\n",
       "        >\n",
       "      >\n",
       "    >\n",
       "  (activation): Swish<\n",
       "    (sigmoid): Sigmoid<>\n",
       "    >\n",
       "  (readout): CellList<\n",
       "    (0): AtomwiseReadout<\n",
       "      (activation): Swish<\n",
       "        (sigmoid): Sigmoid<>\n",
       "        >\n",
       "      (decoder): HalveDecoder<\n",
       "        (activation): Swish<\n",
       "          (sigmoid): Sigmoid<>\n",
       "          >\n",
       "        (output): MLP<\n",
       "          (mlp): SequentialCell<\n",
       "            (0): Dense<\n",
       "              input_channels=128, output_channels=64, has_bias=True, activation=Swish<>\n",
       "              (activation): Swish<\n",
       "                (sigmoid): Sigmoid<>\n",
       "                >\n",
       "              >\n",
       "            (1): Dense<input_channels=64, output_channels=1, has_bias=True>\n",
       "            >\n",
       "          >\n",
       "        >\n",
       "      (aggregator): TensorSummation<>\n",
       "      >\n",
       "    (1): AtomwiseReadout<\n",
       "      (activation): Swish<\n",
       "        (sigmoid): Sigmoid<>\n",
       "        >\n",
       "      (decoder): HalveDecoder<\n",
       "        (activation): Swish<\n",
       "          (sigmoid): Sigmoid<>\n",
       "          >\n",
       "        (output): MLP<\n",
       "          (mlp): SequentialCell<\n",
       "            (0): Dense<\n",
       "              input_channels=128, output_channels=64, has_bias=True, activation=Swish<>\n",
       "              (activation): Swish<\n",
       "                (sigmoid): Sigmoid<>\n",
       "                >\n",
       "              >\n",
       "            (1): Dense<input_channels=64, output_channels=1, has_bias=True>\n",
       "            >\n",
       "          >\n",
       "        >\n",
       "      (aggregator): TensorSummation<>\n",
       "      >\n",
       "    (2): AtomwiseReadout<\n",
       "      (activation): Swish<\n",
       "        (sigmoid): Sigmoid<>\n",
       "        >\n",
       "      (decoder): HalveDecoder<\n",
       "        (activation): Swish<\n",
       "          (sigmoid): Sigmoid<>\n",
       "          >\n",
       "        (output): MLP<\n",
       "          (mlp): SequentialCell<\n",
       "            (0): Dense<\n",
       "              input_channels=128, output_channels=64, has_bias=True, activation=Swish<>\n",
       "              (activation): Swish<\n",
       "                (sigmoid): Sigmoid<>\n",
       "                >\n",
       "              >\n",
       "            (1): Dense<input_channels=64, output_channels=1, has_bias=True>\n",
       "            >\n",
       "          >\n",
       "        >\n",
       "      (aggregator): TensorSummation<>\n",
       "      >\n",
       "    (3): AtomwiseReadout<\n",
       "      (activation): Swish<\n",
       "        (sigmoid): Sigmoid<>\n",
       "        >\n",
       "      (decoder): HalveDecoder<\n",
       "        (activation): Swish<\n",
       "          (sigmoid): Sigmoid<>\n",
       "          >\n",
       "        (output): MLP<\n",
       "          (mlp): SequentialCell<\n",
       "            (0): Dense<\n",
       "              input_channels=128, output_channels=64, has_bias=True, activation=Swish<>\n",
       "              (activation): Swish<\n",
       "                (sigmoid): Sigmoid<>\n",
       "                >\n",
       "              >\n",
       "            (1): Dense<input_channels=64, output_channels=1, has_bias=True>\n",
       "            >\n",
       "          >\n",
       "        >\n",
       "      (aggregator): TensorSummation<>\n",
       "      >\n",
       "    >\n",
       "  (fc_neighbours): FullConnectNeighbours<>\n",
       "  (distances): IndexDistances<\n",
       "    (get_vector): GetVector<>\n",
       "    (norm_last_dim): Norm<axis=-1, keep_dims=False>\n",
       "    >\n",
       "  >"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "net.set_scaleshift([1, 1, 1], 0, readout_id=[0, 1, 3])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "================================================================================\n",
      "Cybertron Engine, Ride-on!\n",
      "--------------------------------------------------------------------------------\n",
      "    Length unit: nm\n",
      "    Input unit scale: 1\n",
      "--------------------------------------------------------------------------------\n",
      "    Deep molecular model:  MolCT\n",
      "--------------------------------------------------------------------------------\n",
      "       Length unit: nm\n",
      "       Atom embedding size: 64\n",
      "       Cutoff distance: 1.0 nm\n",
      "       Radical basis function (RBF): LogGaussianBasis\n",
      "          Minimum distance: 0.04 nm\n",
      "          Maximum distance: 1.0 nm\n",
      "          Reference distance: 1.0 nm\n",
      "          Log Gaussian begin: -3.218876\n",
      "          Log Gaussian end: 0.006724119\n",
      "          Interval for log Gaussian: 0.0512\n",
      "          Sigma for log gaussian: 0.3\n",
      "          Number of basis functions: 64\n",
      "          Rescale the range of RBF to (-1,1).\n",
      "       Calculate distance: Yes\n",
      "       Calculate bond: No\n",
      "       Feature dimension: 128\n",
      "--------------------------------------------------------------------------------\n",
      "       Using 3 independent interaction layers:\n",
      "--------------------------------------------------------------------------------\n",
      "       0. Neural Interaction Unit\n",
      "          Feature dimension: 128\n",
      "          Activation function: Swish\n",
      "          Encoding distance: Yes\n",
      "          Encoding bond: No\n",
      "          Number of heads in multi-haed attention: 8\n",
      "          Use feed forward network: No\n",
      "--------------------------------------------------------------------------------\n",
      "       1. Neural Interaction Unit\n",
      "          Feature dimension: 128\n",
      "          Activation function: Swish\n",
      "          Encoding distance: Yes\n",
      "          Encoding bond: No\n",
      "          Number of heads in multi-haed attention: 8\n",
      "          Use feed forward network: No\n",
      "--------------------------------------------------------------------------------\n",
      "       2. Neural Interaction Unit\n",
      "          Feature dimension: 128\n",
      "          Activation function: Swish\n",
      "          Encoding distance: Yes\n",
      "          Encoding bond: No\n",
      "          Number of heads in multi-haed attention: 8\n",
      "          Use feed forward network: No\n",
      "--------------------------------------------------------------------------------\n",
      "    With 4 readout networks: \n",
      "--------------------------------------------------------------------------------\n",
      "    0. AtomwiseReadout\n",
      "       Activation function: Swish\n",
      "       Decoder: HalveDecoder\n",
      "       Aggregator: TensorSummation\n",
      "       Representation dimension: 128\n",
      "       Readout dimension: 1\n",
      "       Scale: [1.]\n",
      "       Shift: [0.]\n",
      "       Scaleshift mode: Atomwise\n",
      "       Reference value for atom types: None\n",
      "       Output unit: None\n",
      "       Reduce axis: -2\n",
      "--------------------------------------------------------------------------------\n",
      "    1. AtomwiseReadout\n",
      "       Activation function: Swish\n",
      "       Decoder: HalveDecoder\n",
      "       Aggregator: TensorSummation\n",
      "       Representation dimension: 128\n",
      "       Readout dimension: 1\n",
      "       Scale: [1.]\n",
      "       Shift: [0.]\n",
      "       Scaleshift mode: Atomwise\n",
      "       Reference value for atom types: None\n",
      "       Output unit: None\n",
      "       Reduce axis: -2\n",
      "--------------------------------------------------------------------------------\n",
      "    2. AtomwiseReadout\n",
      "       Activation function: Swish\n",
      "       Decoder: HalveDecoder\n",
      "       Aggregator: TensorSummation\n",
      "       Representation dimension: 128\n",
      "       Readout dimension: 1\n",
      "       Scale: 1.0\n",
      "       Shift: 0.0\n",
      "       Scaleshift mode: Atomwise\n",
      "       Reference value for atom types: None\n",
      "       Output unit: None\n",
      "       Reduce axis: -2\n",
      "--------------------------------------------------------------------------------\n",
      "    3. AtomwiseReadout\n",
      "       Activation function: Swish\n",
      "       Decoder: HalveDecoder\n",
      "       Aggregator: TensorSummation\n",
      "       Representation dimension: 128\n",
      "       Readout dimension: 1\n",
      "       Scale: [1.]\n",
      "       Shift: [0.]\n",
      "       Scaleshift mode: Atomwise\n",
      "       Reference value for atom types: None\n",
      "       Output unit: None\n",
      "       Reduce axis: -2\n",
      "--------------------------------------------------------------------------------\n",
      "    Output dimension: [1 1 1 1]\n",
      "    Total output dimension: 4\n",
      "    Output unit for Cybertron: None\n",
      "    Output unit scale: [1. 1. 1. 1.]\n",
      "================================================================================\n"
     ]
    }
   ],
   "source": [
    "net.print_info()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0 model.atom_embedding.embedding_table (64, 128)\n",
      "1 model.dis_filter.linear.weight (128, 64)\n",
      "2 model.dis_filter.linear.bias (128,)\n",
      "3 model.dis_filter.residual.nonlinear.mlp.0.weight (128, 128)\n",
      "4 model.dis_filter.residual.nonlinear.mlp.0.bias (128,)\n",
      "5 model.dis_filter.residual.nonlinear.mlp.1.weight (128, 128)\n",
      "6 model.dis_filter.residual.nonlinear.mlp.1.bias (128,)\n",
      "7 model.interactions.0.positional_embedding.norm.gamma (128,)\n",
      "8 model.interactions.0.positional_embedding.norm.beta (128,)\n",
      "9 model.interactions.0.positional_embedding.x2q.weight (128, 128)\n",
      "10 model.interactions.0.positional_embedding.x2k.weight (128, 128)\n",
      "11 model.interactions.0.positional_embedding.x2v.weight (128, 128)\n",
      "12 model.interactions.0.multi_head_attention.output.weight (128, 128)\n",
      "13 model.interactions.1.positional_embedding.norm.gamma (128,)\n",
      "14 model.interactions.1.positional_embedding.norm.beta (128,)\n",
      "15 model.interactions.1.positional_embedding.x2q.weight (128, 128)\n",
      "16 model.interactions.1.positional_embedding.x2k.weight (128, 128)\n",
      "17 model.interactions.1.positional_embedding.x2v.weight (128, 128)\n",
      "18 model.interactions.1.multi_head_attention.output.weight (128, 128)\n",
      "19 model.interactions.2.positional_embedding.norm.gamma (128,)\n",
      "20 model.interactions.2.positional_embedding.norm.beta (128,)\n",
      "21 model.interactions.2.positional_embedding.x2q.weight (128, 128)\n",
      "22 model.interactions.2.positional_embedding.x2k.weight (128, 128)\n",
      "23 model.interactions.2.positional_embedding.x2v.weight (128, 128)\n",
      "24 model.interactions.2.multi_head_attention.output.weight (128, 128)\n",
      "25 readout.0.decoder.output.mlp.0.weight (64, 128)\n",
      "26 readout.0.decoder.output.mlp.0.bias (64,)\n",
      "27 readout.0.decoder.output.mlp.1.weight (1, 64)\n",
      "28 readout.0.decoder.output.mlp.1.bias (1,)\n",
      "29 readout.1.decoder.output.mlp.0.weight (64, 128)\n",
      "30 readout.1.decoder.output.mlp.0.bias (64,)\n",
      "31 readout.1.decoder.output.mlp.1.weight (1, 64)\n",
      "32 readout.1.decoder.output.mlp.1.bias (1,)\n",
      "33 readout.2.decoder.output.mlp.0.weight (64, 128)\n",
      "34 readout.2.decoder.output.mlp.0.bias (64,)\n",
      "35 readout.2.decoder.output.mlp.1.weight (1, 64)\n",
      "36 readout.2.decoder.output.mlp.1.bias (1,)\n",
      "37 readout.3.decoder.output.mlp.0.weight (64, 128)\n",
      "38 readout.3.decoder.output.mlp.0.bias (64,)\n",
      "39 readout.3.decoder.output.mlp.1.weight (1, 64)\n",
      "40 readout.3.decoder.output.mlp.1.bias (1,)\n",
      "Total parameters:  280196\n"
     ]
    }
   ],
   "source": [
    "tot_params = 0\n",
    "for i, param in enumerate(net.get_parameters()):\n",
    "    tot_params += param.size\n",
    "    print(i, param.name, param.shape)\n",
    "print('Total parameters: ', tot_params)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "n_epoch = 8\n",
    "repeat_time = 1\n",
    "batch_size = 32"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "ds_train = ds.NumpySlicesDataset(\n",
    "    {'R': train_data['R'], 'Z': train_data['Z'], 'E': train_data['E'][:, idx]}, shuffle=True)\n",
    "ds_train = ds_train.batch(batch_size, drop_remainder=True)\n",
    "ds_train = ds_train.repeat(repeat_time)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "ds_valid = ds.NumpySlicesDataset(\n",
    "    {'R': valid_data['R'], 'Z': valid_data['Z'], 'E': valid_data['E'][:, idx]}, shuffle=False)\n",
    "ds_valid = ds_valid.batch(128)\n",
    "ds_valid = ds_valid.repeat(1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "WithLabelLossCell with input type: RZE\n",
      "WithLabelEvalCell with input type: RZE\n",
      "   with scaleshift for training and evaluate dataset:\n",
      "   Output.            Scale           Shift        Mode\n",
      "   0:        1.824854e+01   -4.094204e+02    Atomwise\n",
      "   1:        1.818079e+01   -4.118893e+02    Atomwise\n",
      "   2:        1.816252e+01   -4.142276e+02    Atomwise\n",
      "   3:        1.807273e+01   -3.811549e+02    Atomwise\n",
      "   with reference value for atom types:\n",
      "   Type     Label0    Label1    Label2    Label3\n",
      "   0:        0.00e+00  0.00e+00  0.00e+00  0.00e+00\n",
      "   1:       -1.31e+03 -1.31e+03 -1.31e+03 -1.34e+03\n",
      "   2:        0.00e+00  0.00e+00  0.00e+00  0.00e+00\n",
      "   3:        0.00e+00  0.00e+00  0.00e+00  0.00e+00\n",
      "   4:        0.00e+00  0.00e+00  0.00e+00  0.00e+00\n",
      "   5:        0.00e+00  0.00e+00  0.00e+00  0.00e+00\n",
      "   6:       -9.94e+04 -9.94e+04 -9.94e+04 -9.94e+04\n",
      "   7:       -1.43e+05 -1.43e+05 -1.43e+05 -1.43e+05\n",
      "   8:       -1.97e+05 -1.97e+05 -1.97e+05 -1.97e+05\n",
      "   9:       -2.62e+05 -2.62e+05 -2.62e+05 -2.62e+05\n"
     ]
    }
   ],
   "source": [
    "loss_network = WithLabelLossCell('RZE', net, nn.MAELoss())\n",
    "eval_network = WithLabelEvalCell('RZE', net, nn.MAELoss(), scale=scale, shift=shift, type_ref=ref)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "lr = TransformerLR(learning_rate=1., warmup_steps=4000, dimension=128)\n",
    "optim = nn.Adam(params=net.trainable_params(), learning_rate=lr)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "eval_mae = 'EvalMAE'\n",
    "atom_mae = 'AtomMAE'\n",
    "eval_loss = 'Evalloss'\n",
    "model = Model(loss_network, optimizer=optim, eval_network=eval_network,\n",
    "              metrics={eval_mae: MAE([1, 2], reduce_all_dims=False),\n",
    "                       atom_mae: MAE([1, 2, 3], reduce_all_dims=False, averaged_by_atoms=True),\n",
    "                       eval_loss: MLoss(0)})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "outdir = 'Tutorial_C05'\n",
    "outname = outdir + '_' + net.model_name\n",
    "record_cb = TrainMonitor(model, outname, per_step=16, avg_steps=16,\n",
    "                         directory=outdir, eval_dataset=ds_valid, best_ckpt_metrics=eval_loss)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [],
   "source": [
    "config_ck = CheckpointConfig(save_checkpoint_steps=32, keep_checkpoint_max=64, append_info=[net.hyper_param])\n",
    "ckpoint_cb = ModelCheckpoint(prefix=outname, directory=outdir, config=config_ck)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "[WARNING] ME(7681:140517018826560,MainProcess):2022-08-10-17:17:03.922.843 [mindspore/train/model.py:1097] For TrainMonitor callback, {'step_end', 'begin', 'epoch_end'} methods may not be supported in later version, Use methods prefixed with 'on_train' or 'on_eval' instead when using customized callbacks.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Start training ...\n",
      "Epoch: 1, Step: 16, Learning_rate: 5.2407836e-06, Last_Loss: 27.398788, Avg_loss: 27.347488403320312, EvalMAE: [332.8042  666.40283 413.43115 435.43018], AtomMAE: [20.021822 34.493835 24.720966 22.833675], Evalloss: 25.437191009521484\n",
      "Epoch: 1, Step: 32, Learning_rate: 1.0830951e-05, Last_Loss: 18.250584, Avg_loss: 23.159127116203308, EvalMAE: [252.40771 539.66016 314.9839  328.2124 ], AtomMAE: [14.946562 27.936647 18.9759   17.622032], Evalloss: 19.754390716552734\n",
      "Epoch: 2, Step: 48, Learning_rate: 1.6421121e-05, Last_Loss: 14.260176, Avg_loss: 17.202343702316284, EvalMAE: [236.59326 351.68652 252.01074 260.43994], AtomMAE: [13.32491  18.653118 14.712646 14.975773], Evalloss: 15.148639678955078\n",
      "Epoch: 2, Step: 64, Learning_rate: 2.2011289e-05, Last_Loss: 14.724831, Avg_loss: 13.160118043422699, EvalMAE: [223.96143 267.35254 254.60645 249.40186], AtomMAE: [12.666828 15.647819 14.348164 14.614561], Evalloss: 13.69890022277832\n",
      "Epoch: 3, Step: 80, Learning_rate: 2.760146e-05, Last_Loss: 9.907473, Avg_loss: 12.792615234851837, EvalMAE: [211.40625 256.78076 244.8628  235.18799], AtomMAE: [11.805235 15.073635 13.855062 13.590918], Evalloss: 13.050618171691895\n",
      "Epoch: 3, Step: 96, Learning_rate: 3.3191627e-05, Last_Loss: 11.705648, Avg_loss: 11.857816874980927, EvalMAE: [190.40967 250.58691 234.04639 219.5664 ], AtomMAE: [10.716079 14.008486 13.251503 12.608332], Evalloss: 12.313039779663086\n",
      "Epoch: 4, Step: 112, Learning_rate: 3.8781796e-05, Last_Loss: 11.893072, Avg_loss: 10.675153017044067, EvalMAE: [168.71777 226.25    217.24268 198.21631], AtomMAE: [ 9.535676 13.144618 12.57393  11.562946], Evalloss: 11.154577255249023\n",
      "Epoch: 4, Step: 128, Learning_rate: 4.437197e-05, Last_Loss: 10.478094, Avg_loss: 10.656313180923462, EvalMAE: [154.49854 219.396   220.24756 184.63574], AtomMAE: [ 8.525218  12.013991  12.094766  10.2883625], Evalloss: 10.719141006469727\n",
      "Epoch: 5, Step: 144, Learning_rate: 4.9962135e-05, Last_Loss: 8.520898, Avg_loss: 9.38396018743515, EvalMAE: [122.94629 181.9126  179.74463 152.04883], AtomMAE: [ 6.9292574 10.348976  10.296545   8.585039 ], Evalloss: 8.763139724731445\n",
      "Epoch: 5, Step: 160, Learning_rate: 5.5552304e-05, Last_Loss: 7.22596, Avg_loss: 7.9867982268333435, EvalMAE: [113.214355 159.4541   159.36328  125.26709 ], AtomMAE: [6.5375104 8.864579  8.98425   7.198041 ], Evalloss: 7.669915676116943\n",
      "Epoch: 6, Step: 176, Learning_rate: 6.114247e-05, Last_Loss: 6.9385834, Avg_loss: 7.152495086193085, EvalMAE: [127.81885 135.48877 131.50293 118.21631], AtomMAE: [6.99376   7.536427  7.468955  6.6130886], Evalloss: 7.059635162353516\n",
      "Epoch: 6, Step: 192, Learning_rate: 6.6732646e-05, Last_Loss: 6.029707, Avg_loss: 6.011962831020355, EvalMAE: [ 85.30029  110.489746 104.58301   92.28027 ], AtomMAE: [4.8831496 6.22881   5.951734  5.343926 ], Evalloss: 5.40406608581543\n",
      "Epoch: 7, Step: 208, Learning_rate: 7.2322815e-05, Last_Loss: 5.7389402, Avg_loss: 5.471996396780014, EvalMAE: [87.418945 93.98291  85.884766 87.39111 ], AtomMAE: [5.007127  5.34113   4.9300346 5.017578 ], Evalloss: 4.881006240844727\n",
      "Epoch: 7, Step: 224, Learning_rate: 7.791298e-05, Last_Loss: 4.6558332, Avg_loss: 4.634718626737595, EvalMAE: [88.07764  87.01416  79.35742  83.927734], AtomMAE: [5.093235 4.966573 4.57263  4.851916], Evalloss: 4.6562418937683105\n",
      "Epoch: 8, Step: 240, Learning_rate: 8.3503146e-05, Last_Loss: 5.6020756, Avg_loss: 5.011868670582771, EvalMAE: [87.75    88.51123 83.05127 85.15332], AtomMAE: [5.0828333 5.1343884 4.800145  4.9472466], Evalloss: 4.740358829498291\n",
      "Epoch: 8, Step: 256, Learning_rate: 8.9093315e-05, Last_Loss: 4.3972125, Avg_loss: 4.2410566210746765, EvalMAE: [ 97.82568   90.83008   91.20459  102.581055], AtomMAE: [5.4338207 5.124866  5.135529  5.7657886], Evalloss: 5.263698577880859\n",
      "Training Fininshed!\n",
      "Training Time: 00:00:39\n"
     ]
    }
   ],
   "source": [
    "print(\"Start training ...\")\n",
    "beg_time = time.time()\n",
    "model.train(n_epoch, ds_train, callbacks=[record_cb, ckpoint_cb], dataset_sink_mode=False)\n",
    "end_time = time.time()\n",
    "used_time = end_time - beg_time\n",
    "m, s = divmod(used_time, 60)\n",
    "h, m = divmod(m, 60)\n",
    "print(\"Training Fininshed!\")\n",
    "print(\"Training Time: %02d:%02d:%02d\" % (h, m, s))"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3.7.5 ('mindsponge')",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.5"
  },
  "orig_nbformat": 4,
  "vscode": {
   "interpreter": {
    "hash": "2496ecc683137a232cae2452fbbdd53dab340598b6e499c8995be760f3a431b4"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
